import os
import cv2
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 读取预测结果
subm = pd.read_csv('./subtt_b_2.csv', sep='\t', header=None, names=['name', 'mask'])

# 读取测试数据
test_mask = pd.read_csv('./data/test_b_samplesubmit.csv', sep='\t', names=['name', 'mask'])
test_mask['name'] = test_mask['name'].apply(lambda x: './data/train/' + x)

def rle_decode(mask_rle, shape=(512, 512)):
    '''
    mask_rle: run-length as string formated (start length)
    shape: (height,width) of array to return
    Returns numpy array, 1 - mask, 0 - background
    '''
    s = mask_rle.split()
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape, order='F')

# 显示对比图
for idx, row in subm.iterrows():
    if idx >= 20:
        break
    image_path = test_mask.loc[test_mask['name'].str.contains(row['name'])]['name'].values[0]
    image = cv2.imread(image_path)
    pred_mask = rle_decode(row['mask'])

    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.title('Original Image')

    plt.subplot(1, 2, 2)
    plt.imshow(pred_mask, cmap='gray')
    plt.title('Predicted Mask')

    plt.show()